[OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel - #2743

…out the session

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The patch and its module flags sit in model_load_utils with the
expert-parallel loader that applies them, so the test drives and resets
them there.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The patch fuses the Fp8Dequantize that transformers prepends to the
routed-expert converter into its MergeModulelist, so the loader still
gives each expert its tensor_idx. Only the install was tested. The new
test runs the patched update_weight_conversions on an expert converter
and the resulting chain on per-expert block-FP8 weights. It fails if the
converter no longer leads with a MergeModulelist, or if the fused op
skips the dequant or the merge. transformers 5.18's chain passes an
output_dtype the patched dequant does not take, so the chain runs that
version's own dequant; the patched arithmetic has its own tests.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

- The two block-size cases, a grid finer than the configured block and
  no block configured, become rows of the blockwise-reference table, which
  gains configured and grid columns.
- The dtype test keeps its unrequested (bf16) and requested (fp16) cases;
  requesting bf16 or fp32 took the same branch.
- The version test keeps the 5.10.1 boundary; 5.17.0 took the same branch.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The pre-5.10.1 dequant patch emitted a non-empty weight unscaled when its scale
grid was empty. That only happens when a weight and its scales were sliced
apart, as transformers 5.9's expert tensor_idx bug did on expert-parallel loads,
so it loaded wrong weights without a word; upstream's own dequant fails on
such a grid. It also fell back to the configured block for a grid neither that
block nor an even split gives, scaling rows wrongly. Raise in both cases; an
empty weight still passes through its empty grid.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

From 5.10.1 the patch kept transformers' own dequant untouched, on the
assumption that no shipping model has a partial last block there.
DeepSeek-V3, V3.1 and Kimi-K2-Instruct do: kv_a_proj_with_mqa has 576
rows under a 5-row grid of 128-row blocks, and Fp8Dequantize still
raises "not divisible by scale grid" for it, so their block-FP8 loads
failed. The stock dequant is kept, for its ue8m0 and FP4 decoding and
output_dtype, and wrapped: an FP8 weight whose grid is the configured
block's ceil grid is padded to whole blocks for it and cropped back.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

When fusing transformers' prepended expert dequant into its merge
failed, the patch printed a warning and went on. Every routed expert
then loaded sliced apart from its scales. The failure now raises. A
transformers without the update_weight_conversions hook prepends no
dequant, so there is nothing to fuse and the patch skips it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…tizes

The load dequantizes native block-FP8 checkpoints as it loads, so they
join the formats it accepts.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…heck

The hub-ID test now loads an unquantized, a block-FP8 and an MXFP4
config: each must reach the checkpoint's resolution, past the format
check, so dropping either format from the accepted set fails it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

From transformers 5.10.1 the expert-parallel load wraps transformers' block-FP8 dequant, which derives
the block size from the scale grid, and padded a weight to whole blocks only when it did not divide
into that grid. A partial last block can still divide evenly: 448 rows under a 4-row grid of
128-row blocks go to the dequant as blocks of 112, and every row group after the first gets the
wrong scale, silently. The wrapper now pads whenever the weight is not block-aligned, which is a
no-op for whole blocks.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The pre-5.10.1 dequant multiplies its own fp32 copy of the weight by the tiled scales in place,
so a weight holds two full-size fp32 temporaries instead of three while it loads. The dequant is
rebound after the expert converter's fuse, so a refused patch leaves transformers as it was.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The rule that a block-FP8 checkpoint's dense part is not sharded lands here, with the load that
makes such a checkpoint loadable, rather than in [10/21], where it preempted the format refusal.
Left to decide, the plan replicates that dense part and says why; asked to shard it, the decision
refuses, naming MODELOPT_FSDP_SIZE=1, and the loader refuses a plan that shards it.

The pre-5.10.1 dequant scales its fp32 copy of a weight in place, through a view of its whole
column blocks and then the partial last one, against scales tiled over the rows only. It builds no
full-size fp32 scale tile: an (18432, 7168) weight peaks at 6 bytes per element instead of 8. A
weight narrower than one block joins the reference table.

The expert converter's fuse no longer skips a merge it already fused, a layout transformers'
update_weight_conversions never produces (it builds fresh converters on every call); fusing it
again would dequantize once all the same. The load's log names the dtype the dequant emits.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

quant_method "fp8" is not block FP8 alone: per-tensor FP8 checkpoints (AutoFP8's: a scalar
weight_scale per weight, no weight_block_size) name the same method. The expert-parallel load
admitted them, and transformers' dequant reads them as 128x128 blocks with no weight_scale_inv, so
their weights load unscaled without an error (on transformers 5.18 a tiny Llama loaded that way
has max |w| 448 where the original has 0.079). _ckpt_quant_method now calls an fp8 checkpoint
without a weight_block_size "fp8-per-tensor", which the load refuses, before anything loads, as a
format it cannot dequantize; the decision's block-FP8 rule no longer applies to it either.

The decision, the loader's refusal and its docstring give the reason a block-FP8 dense part cannot
be sharded as Distributed Checkpoint reading it as stored, with no dequantization: its scales'
format was never the blocker (DeepSeek-V3's are fp32). A comment says where the dequant's block size
comes from: transformers keeps the checkpoint's quantization_config and takes only the loading
attributes from the one the load passes.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A config the load is handed can name its dtype, or carry "auto" for the checkpoint's own, and the
dequant's Tensor.to reads a string as a device. The dtype a string names is now taken as the torch
dtype, and "auto" as the bfloat16 default the dequant emits unrequested.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…tize

Every transformers the repo supports ships Fp8Dequantize (5.5 does), so the skip never fires, and
the load dereferences the class on the same releases.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The override installs below transformers 5.10.1 and took exactly (quantized, scales): 5.8.0 to
5.10.0 call it that way, and 5.10.1, where the gate switches to transformers' own dequant, is the
first to pass output_dtype. A caller passing it would have stopped the load with a bare TypeError.
It now takes output_dtype and, when given, emits it: the loader passes the destination dtype, which
the module-level model dtype only reconstructs.

The expert-converter test runs the patched dequant through this transformers' chain, which passes
output_dtype, instead of restoring the stock one around it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…uant

From transformers 5.10.1 a partial last block is padded to whole blocks for transformers' own
dequant. The padding cast the weight to fp32 first, so the dequant got an input dtype it never sees
from the loader, and the load an fp32 copy it does not need. The weight is now padded through a
same-size uint8 view, in its own fp8 dtype; the zero byte is zero in every fp8 format.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

From transformers 5.10.1 the patch wraps transformers' dequant, calling the method it replaces, and
only this module's flag kept it from wrapping again. Re-importing the module loses that flag, so the
wrapper wrapped itself, one frame per re-import. The dequant op now carries a flag of its own, as
the FP8 quantizer does for its converter hook.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The expert-parallel load passes FineGrainedFP8Config(dequantize=True), and both dequant branches read
weight_block_size off the quantizer's config, relying on transformers to keep the checkpoint's
quantization_config and take only the loading attributes from the one passed. A test now pins that
merge on whichever transformers runs the suite.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Its reason, that Distributed Checkpoint reads block FP8 as stored, holds at any EP degree, but its
remedy, replicating the dense part to dequantize it as it loads, holds only where the
expert-parallel load runs. Without expert parallelism the FSDP2 load refuses a quantized checkpoint
itself.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…ckpoint

An fp8 checkpoint without weight_block_size was told the load "cannot dequantize a 'fp8-per-tensor'
checkpoint" while listing fp8 among what it dequantizes, and the refusal never named the block size
that tells the two apart. It now says the load dequantizes block-fp8 checkpoints, which name their
block in quantization_config.weight_block_size, and that this one names none.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A per-tensor FP8 checkpoint names no weight_block_size, which is how the load tells it from block
FP8 and refuses it. transformers' FP8 config object fills in a 128x128 block whatever the
checkpoint names, so a per-tensor checkpoint whose config came in as that object, as a loaded
model's does, read as block FP8 and would load unscaled. For an FP8 config object, the block now
comes from the checkpoint's own config.json, where both the decision and the load can read it.

The plan decision also stops offering an EP degree for a block-fp8 dense part sharded within
several expert-data-parallel replicas: no EP degree lets that dense part shard.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Both patches unpacked the last two dims first, so a weight or scale grid of fewer than two dims
failed with a bare unpack error from inside the patched method. From 5.10.1 such a call now goes to
transformers' own dequant, as every case the patch does not place does; below it, the override
refuses it by name. The override also drops its FP4 branch: it called an unpack transformers only
added in 5.10.1, so it could not run on the releases the override replaces.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The per-tensor check reads an FP8 config object's block from the checkpoint's own config.json, but
a path without one, a Hub ID the load has not made local yet, left the object's default 128x128
block standing, and a per-tensor checkpoint would still load unscaled. Such a config is now
"fp8-unverified", which the load refuses, naming the two ways in: the config as
AutoConfig.from_pretrained reads it, or a local checkpoint. Reading the file decides it on every
rank alike, without fetching anything.

For a block-fp8 dense part within several expert-data-parallel replicas, the remedy also names
starting from an unquantized copy, which on the warning path is the one that changes the outcome.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>